import numpy as np
from light_training.dataloading.dataset import get_loader
import torch 
import torch.nn as nn 
from monai.networks.nets.basic_unet import BasicUNet
from monai.networks.nets.unetr import UNETR
from monai.networks.nets.swin_unetr import SwinUNETR
from monai.inferers import SlidingWindowInferer
from light_training.evaluation.metric import dice
from light_training.trainer import Trainer
from monai.utils import set_determinism
from light_training.utils.files_helper import save_new_model_and_delete_last
from models.uent3d import UNet3D
from monai.networks.nets.segresnet import SegResNet
from models.transbts.TransBTS_downsample8x_skipconnection import TransBTS
from einops import rearrange
from models.modelgenesis.unet3d import UNet3DModelGen
from models.transvw.models.ynet3d import UNet3DTransVW
from monai.networks.nets.basic_unet import BasicUNet
from monai.networks.nets.attentionunet import AttentionUnet
from light_training.loss.compound_losses import DC_and_CE_loss
from light_training.loss.dice import MemoryEfficientSoftDiceLoss
from light_training.evaluation.metric import dice
set_determinism(123)
from light_training.loss.compound_losses import DC_and_CE_loss
import os
from medpy import metric

# os.environ["CUDA_VISIBLE_DEVICES"] = "1,2,3"
data_dir = "./data/fullres/train"

fold = 0
logdir = f"./logs/nnunetsetings_lr1ef2ep1000_ce_dc_fold_{fold}"

env = "pytorch"
model_save_path = os.path.join(logdir, "model")
max_epoch = 1000
batch_size = 2
val_every = 2
num_gpus = 1
device = "cuda:2"
patch_size = [128, 128, 128]

class BraTSTrainer(Trainer):
    def __init__(self, env_type, max_epochs, batch_size, device="cpu", val_every=1, num_gpus=1, logdir="./logs/", master_ip='localhost', master_port=17750, training_script="train.py"):
        super().__init__(env_type, max_epochs, batch_size, device, val_every, num_gpus, logdir, master_ip, master_port, training_script)
        self.window_infer = SlidingWindowInferer(roi_size=patch_size,
                                        sw_batch_size=2,
                                        overlap=0.6)
        self.patch_size = patch_size
        self.all_labels=[1, 2, 3]

        from models.nnunet3d import get_nnunet3d
        self.model = get_nnunet3d(4, 4)

        # self.model = SwinUNETR(patch_size, 1, 1, feature_size=48)
        # self.model = UNet2D()
        # self.model = BasicUNet(3, 1, 17)
        # self.model = UNETR(1, 17, [96, 96, 96])
        # _, model = TransBTS(dataset='brats', _conv_repr=True, _pe_type="learned")
        # self.model = model

        # self.model = SegResNet(3, 32, 1, 17)

        self.best_mean_dice = 0.0
        # self.optimizer = torch.optim.AdamW(self.model.parameters(), lr=1e-2, weight_decay=3e-5)
        # self.optimizer = torch.optim.AdamW(self.model.parameters(), lr=1e-2, weight_decay=3e-5)
        self.optimizer = torch.optim.SGD(self.model.parameters(), lr=1e-2, weight_decay=3e-5,
                                    momentum=0.99, nesterov=True)
        
        self.scheduler_type = "poly"
        # self.scheduler_type = "cosine_with_warmup"
        # self.warmup = 0.03

        self.loss_func = DC_and_CE_loss({'batch_dice': True,
                                   'smooth': 1e-5, 'do_bg': False, 'ddp': False}, {}, weight_ce=1, weight_dice=1,
                                  ignore_label=None, dice_class=MemoryEfficientSoftDiceLoss)
        
        # self.loss_func = nn.CrossEntropyLoss()
    
    def training_step(self, batch):
        import time 
        image, label = self.get_input(batch)

        pred = self.model(image)
        loss = self.loss_func(pred, label)

        self.log("train_loss", loss, step=self.global_step)
        return loss 

    # for image, label in data_loader:
    def get_input(self, batch):
        image = batch["data"]
        label = batch["seg"]

        label = label.float()
        # label = label[:, 0].long()

        return image, label 

    def cal_metric(self, gt, pred, voxel_spacing=[1.0, 1.0, 1.0]):
        if pred.sum() > 0 and gt.sum() > 0:
            d = dice(pred, gt)
            # hd95 = metric.binary.hd95(pred, gt)
            return np.array([d, 50])
        
        elif gt.sum() == 0 and pred.sum() == 0:
            return np.array([1.0, 50])
        
        else:
            return np.array([0.0, 50])
    
    def validation_step(self, batch):
        image, label = self.get_input(batch)
       
        output = self.model(image).argmax(dim=1).cpu().numpy()
        # output = self.window_infer(image, self.model).argmax(dim=1).cpu().numpy()
        target = label.cpu().numpy()
        if len(target.shape) == 5:
            target = target[:, 0]
        
        dices = []

        c = 4
        for i in range(1, c):
            pred_c = (output == i)
            target_c = (target == i)
            cal_dice, cal_hd = self.cal_metric(target_c, pred_c)
            dices.append(cal_dice)
        
        return dices
    
    def validation_end(self, mean_val_outputs, val_outputs):
        dices = mean_val_outputs
        
        print(f"dices is {dices}")

        mean_dice = sum(dices) / len(dices)
        
        index = 0
        for d in dices:
            index += 1
            self.log(f"dice_{index}", d, step=self.epoch)

        self.log("mean_dice", mean_dice, step=self.epoch)

        if mean_dice > self.best_mean_dice:
            self.best_mean_dice = mean_dice
            save_new_model_and_delete_last(self.model, 
                                            os.path.join(model_save_path, 
                                            f"best_model_{mean_dice:.4f}.pt"), 
                                            delete_symbol="best_model")

        save_new_model_and_delete_last(self.model, 
                                        os.path.join(model_save_path, 
                                        f"final_model_{mean_dice:.4f}.pt"), 
                                        delete_symbol="final_model")

        print(f"mean_dice is {mean_dice}")

if __name__ == "__main__":

    trainer = BraTSTrainer(env_type=env,
                            max_epochs=max_epoch,
                            batch_size=batch_size,
                            device=device,
                            logdir=logdir,
                            val_every=val_every,
                            num_gpus=num_gpus,
                            master_port=17751,
                            training_script=__file__)
    
    train_ds, val_ds, test_ds = get_loader(data_dir=data_dir, fold=fold)

    trainer.train(train_dataset=train_ds, val_dataset=val_ds)
